#pragma once
// clang-format off
#include <Python.h>
// clang-format on
#include <torch/torch.h>

#include <string_view>
#include <tuple>

#include "dose_torch_export.h"

extern "C"
{
    DOSE_TORCH_EXPORT PyObject* PyInit_dose_torch(void);
}

namespace dose_torch
{

    // 计算DVH(剂量体积直方图)的PyTorch自定义算子
    // 输入: voxelVolume, voxelDose, doExactAccumulation
    // 输出: (binned dose values, accumulated relative volumes)
    DOSE_TORCH_EXPORT std::tuple<torch::Tensor, torch::Tensor> accumulate(const torch::Tensor& voxelVolume,
                                                                          const torch::Tensor& voxelDose,
                                                                          bool doExactAccumulation);

    // 获取特定体积百分比对应的剂量的PyTorch自定义算子
    DOSE_TORCH_EXPORT torch::Scalar get_metric(const torch::Tensor& binnedDoseValues,
                                               const torch::Tensor& accumulatedRelativeVolumes,
                                               double inDoseAtVolume);

    // 计算ROI内的剂量统计指标的PyTorch自定义算子
    DOSE_TORCH_EXPORT torch::Scalar get_dose_stat_roi_interp(const torch::Tensor& binnedDoseValues,
                                                             const torch::Tensor& accumulatedRelativeVolumes,
                                                             const torch::Tensor& voxelVolume,
                                                             const torch::Tensor& voxelDose,
                                                             std::string_view dose_type,
                                                             double doseAtVolume = 0.0);

}  // namespace dose_torch